"""
LLM Agent客户端

提供LLM Agent调用功能，支持字段提取和字段溯源
"""

import re
import json
import httpx
from typing import Dict, Any, Optional, List, Union
from loguru import logger

from config.llm_config import AgentConfig
from src.services.llm.models import (
    AgentRequest,
    AgentResponse,
    ExtractionState,
    TracingState,
    AuditState,
    ClassificationState,
    ExtractionResult,
    TracingResult,
    AuditResult,
    ClassificationResult
)


class LLMClient:
    """LLM Agent基础客户端"""
    
    def __init__(self, config: AgentConfig):
        """
        初始化LLM客户端
        
        Args:
            config: Agent配置
        """
        self.config = config
        self.client = httpx.AsyncClient(
            timeout=config.timeout,
            headers={
                "Authorization": f"Bearer {config.api_key}",
                "Content-Type": "application/json"
            }
        )
    
    async def call_agent(
        self,
        state: Dict[str, Any],
        user_input: str = ""
    ) -> AgentResponse:
        """
        调用Agent
        
        Args:
            state: 状态参数
            user_input: 用户输入
            
        Returns:
            AgentResponse: Agent响应
            
        Raises:
            Exception: 调用失败时抛出异常
        """
        request = AgentRequest(
            agentId=self.config.agent_id,
            userChatInput=user_input,
            state=state
        )
        
        max_retries = self.config.max_retries
        last_error = None
        
        for attempt in range(max_retries):
            try:
                logger.debug(f"调用Agent，尝试 {attempt + 1}/{max_retries}")
                
                response = await self.client.post(
                    self.config.base_url,
                    json=request.model_dump()
                )
                response.raise_for_status()
                
                response_data = response.json()
                logger.debug(f"响应数据: {response_data}")
                
                return AgentResponse(**response_data)
                
            except httpx.HTTPStatusError as e:
                last_error = e
                logger.warning(f"Agent调用HTTP错误 (尝试 {attempt + 1}/{max_retries}): {e}")
                logger.warning(f"响应内容: {e.response.text if e.response else 'N/A'}")
                
            except Exception as e:
                last_error = e
                logger.warning(f"Agent调用异常 (尝试 {attempt + 1}/{max_retries}): {e}")
        
        # 所有重试都失败
        error_msg = f"Agent调用失败，已重试{max_retries}次: {last_error}"
        logger.error(error_msg)
        raise Exception(error_msg)
    
    async def close(self):
        """关闭客户端"""
        await self.client.aclose()
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器退出"""
        await self.close()


class ExtractionAgent:
    """字段提取Agent"""
    
    def __init__(self, client: LLMClient):
        """
        初始化提取Agent
        
        Args:
            client: LLM客户端
        """
        self.client = client
    
    async def extract_field(
        self,
        file_content: str,
        field_name: str,
        extraction_logic: str,
        reference_materials: str
    ) -> ExtractionResult:
        """
        提取字段
        
        Args:
            file_content: 文件内容（content_list.json内容）
            field_name: 字段名称
            extraction_logic: 提取逻辑
            reference_materials: 参考材料
            
        Returns:
            ExtractionResult: 提取结果
        """
        try:
            state = ExtractionState(
                file_content=file_content,
                field_name=field_name,
                extraction_logic=extraction_logic,
                reference_materials=reference_materials
            )
            
            state_dict = state.model_dump()
            logger.info(f"调用提取Agent - 字段: {field_name}")
            
            response = await self.client.call_agent(state=state_dict)
            
            # 解析提取结果，取最后一个choice
            if not response.choices:
                logger.warning(f"提取Agent响应中没有choices")
                return ExtractionResult(
                    raw_response="",
                    success=False,
                    error="响应中没有choices"
                )
            
            last_choice = response.choices[-1].content
            logger.info(f"提取Agent响应内容: {last_choice}")
            
            # 使用正则解析JSON对象
            extraction_data = self._parse_extraction_json(last_choice)
            
            extracted_value = extraction_data.get('value') if extraction_data else None
            confidence = extraction_data.get('confidence') if extraction_data else None
            
            if extracted_value is None:
                logger.warning(f"无法从响应中解析提取结果，完整响应: {last_choice}")
            
            return ExtractionResult(
                extracted_value=extracted_value,
                confidence=confidence,
                raw_response=last_choice,
                success=extracted_value is not None,
                error=None if extracted_value else "无法解析提取结果"
            )
            
        except Exception as e:
            logger.error(f"字段提取失败: {e}", exc_info=True)
            return ExtractionResult(
                raw_response="",
                success=False,
                error=str(e)
            )
    
    def _parse_extraction_json(self, content: str) -> Optional[Dict[str, Any]]:
        """
        解析提取结果的 JSON 数据
        
        从响应内容中提取JSON对象
        
        Args:
            content: 响应内容
            
        Returns:
            Optional[Dict]: 包含 value 和 confidence 的字典，解析失败返回None
        """
        try:
            # 尝试匹配 ```json ... ``` 格式
            json_pattern = r'```json\s*(\{.*?\})\s*```'
            match = re.search(json_pattern, content, re.DOTALL)
            
            data = None
            if match:
                json_str = match.group(1)
                data = json.loads(json_str)
            else:
                # 如果没有匹配到，尝试直接解析整个内容
                try:
                    data = json.loads(content)
                except:
                    pass
            
            if data:
                # 提取字段值
                extracted_value = None
                for key in ["提取结果", "extracted_value", "result", "value"]:
                    if key in data:
                        extracted_value = str(data[key])
                        break
                
                # 提取 confidence，如果不存在则使用默认值 0.8
                confidence = data.get("confidence") or data.get("置信度") or 0.8
                
                return {
                    'value': extracted_value,
                    'confidence': confidence
                }
            
            logger.warning(f"无法解析提取结果，原始内容: {content[:200]}...")
            return None
            
        except Exception as e:
            logger.error(f"解析提取结果异常: {e}")
            return None


class TracingAgent:
    """字段溯源Agent"""
    
    def __init__(self, client: LLMClient):
        """
        初始化溯源Agent
        
        Args:
            client: LLM客户端
        """
        self.client = client
    
    async def trace_field(
        self,
        file_content: str,
        field_name: str,
        extraction_logic: str,
        reference_materials: str,
        extraction_result: str
    ) -> TracingResult:
        """
        溯源字段
        
        Args:
            file_content: 文件内容（middle.json内容）
            field_name: 字段名称
            extraction_logic: 提取逻辑
            reference_materials: 参考材料
            extraction_result: 提取结果
            
        Returns:
            TracingResult: 溯源结果
        """
        try:
            state = TracingState(
                file_content=file_content,
                field_name=field_name,
                extraction_logic=extraction_logic,
                reference_materials=reference_materials,
                extraction_result=extraction_result
            )
            
            state_dict = state.model_dump()
            logger.info(f"调用溯源Agent - 字段: {field_name}")
            
            response = await self.client.call_agent(state=state_dict)
            
            # 解析溯源结果，取最后一个choice
            if not response.choices:
                logger.warning(f"溯源Agent响应中没有choices")
                return TracingResult(
                    raw_response="",
                    success=False,
                    error="响应中没有choices"
                )
            
            last_choice = response.choices[-1].content
            logger.info(f"溯源Agent响应内容: {last_choice}")
            
            # 解析溯源结果（V2 - 支持多 regions）
            tracing_data = self._parse_tracing_result(last_choice)
            
            if tracing_data and isinstance(tracing_data, dict):
                # 检查状态
                status = tracing_data.get("status", "success")
                if status == "not_found":
                    return TracingResult(
                        raw_response=last_choice,
                        success=False,
                        error="LLM未找到匹配位置"
                    )
                
                # 解析 regions（新格式）
                regions_data = tracing_data.get("regions", [])
                regions = []
                
                if regions_data and isinstance(regions_data, list):
                    from src.services.llm.models import Region
                    for region_dict in regions_data:
                        if not isinstance(region_dict, dict):
                            continue
                        page_idx = region_dict.get("page_idx")
                        bbox = region_dict.get("bbox")
                        if page_idx is not None and bbox:
                            regions.append(Region(
                                page_idx=page_idx,
                                page_number=page_idx + 1,
                                bbox=bbox,
                                text_segment=None
                            ))
                
                # 兼容旧格式（单个 bbox）
                if not regions:
                    page_idx = tracing_data.get("page_idx")
                    bbox = tracing_data.get("bbox")
                    if page_idx is not None and bbox:
                        from src.services.llm.models import Region
                        regions.append(Region(
                            page_idx=page_idx,
                            page_number=page_idx + 1,
                            bbox=bbox,
                            text_segment=None
                        ))
                
                if not regions:
                    return TracingResult(
                        raw_response=last_choice,
                        success=False,
                        error="LLM返回的坐标数据无效"
                    )
                
                # 构建返回结果（兼容旧字段）
                first_region = regions[0]
                return TracingResult(
                    matched_text=extraction_result,
                    regions=regions,
                    page_idx=first_region.page_idx,
                    page_number=first_region.page_number,
                    bbox=first_region.bbox,
                    match_type="llm",
                    raw_response=last_choice,
                    success=True,
                    error=None
                )
                
            else:
                error_msg = "无法解析溯源结果"
                if tracing_data and not isinstance(tracing_data, dict):
                    error_msg = f"溯源结果类型错误: {type(tracing_data)}"
                logger.warning(f"{error_msg}, 原始响应: {last_choice[:200]}...")
                return TracingResult(
                    raw_response=last_choice,
                    success=False,
                    error=error_msg
                )
            
        except Exception as e:
            logger.error(f"字段溯源失败: {e}", exc_info=True)
            return TracingResult(
                raw_response="",
                success=False,
                error=str(e)
            )
    
    def _normalize_matched_text(self, matched_text: Any) -> Optional[str]:
        """
        标准化 matched_text 为字符串
        
        处理多种类型：
        - str: 直接返回
        - dict/list: 转为 JSON 字符串
        - None: 返回 None
        - 其他: 转为字符串
        
        Args:
            matched_text: 原始 matched_text（可能是 str/dict/list/None）
            
        Returns:
            Optional[str]: 标准化后的字符串
        """
        if matched_text is None:
            return None
        
        if isinstance(matched_text, str):
            return matched_text
        
        if isinstance(matched_text, (dict, list)):
            try:
                return json.dumps(matched_text, ensure_ascii=False)
            except Exception as e:
                logger.warning(f"matched_text 转 JSON 失败: {e}, 使用 str() 转换")
                return str(matched_text)
        
        # 其他类型，直接转字符串
        return str(matched_text)
    
    def _parse_tracing_result(self, content: str) -> Optional[Dict[str, Any]]:
        """
        解析溯源结果
        
        从响应内容中提取JSON对象中的定位结果
        
        Args:
            content: 响应内容
            
        Returns:
            Optional[Dict]: 溯源数据，解析失败返回None
        """
        try:
            # 尝试匹配 ```json ... ``` 格式
            json_pattern = r'```json\s*(\{.*?\})\s*```'
            match = re.search(json_pattern, content, re.DOTALL)
            
            if match:
                json_str = match.group(1)
                data = json.loads(json_str)
                
                logger.debug(f"解析到的JSON数据类型: {type(data)}, 内容: {data}")
                
                # 尝试获取定位结果
                result = None
                if "定位结果" in data:
                    result = data["定位结果"]
                elif "tracing_result" in data:
                    result = data["tracing_result"]
                elif "location" in data:
                    result = data["location"]
                
                # 类型检查：确保返回的是字典
                if result is not None:
                    if isinstance(result, dict):
                        logger.debug(f"溯源结果解析成功: {result}")
                        return result
                    else:
                        logger.error(f"溯源结果类型错误: 期望dict，实际为 {type(result)}, 内容: {result}")
                        return None
            
            logger.warning(f"无法解析溯源结果，原始内容: {content[:200]}...")
            return None
            
        except Exception as e:
            logger.error(f"解析溯源结果异常: {e}", exc_info=True)
            return None


class AuditAgent:
    """字段审核Agent"""
    
    def __init__(self, client: LLMClient):
        """
        初始化审核Agent
        
        Args:
            client: LLM客户端
        """
        self.client = client
    
    async def audit_field(
        self,
        audit_type: str,
        field_name: str,
        audit_content: str,
        auto_process_rules: str,
        audit_data_sources: Dict[str, Any],
        user_input_value: Optional[str],
        reference_materials: Dict[str, Any]
    ) -> AuditResult:
        """
        审核字段
        
        Args:
            audit_type: 审核类型（字段审核/文件审核）
            field_name: 字段名称
            audit_content: 审核内容
            auto_process_rules: 自动处理规则
            audit_data_sources: 审核数据源（结构化对象）
            user_input_value: 用户输入的字段值
            reference_materials: 参考材料（从附件提取的结果，结构化对象）
            
        Returns:
            AuditResult: 审核结果
        """
        try:
            # 构建审核数据源字符串（格式化为易读的文本）
            data_sources_str = self._format_audit_data_sources(audit_data_sources)
            
            # 构建用户输入键值对
            user_input_key = f"{field_name}：{user_input_value}" if user_input_value else ""
            
            # 构建参考材料字符串（格式化为易读的文本）
            reference_materials_str = self._format_reference_materials(reference_materials)
            
            state = AuditState(
                audit_type=audit_type,
                field_name=field_name,
                audit_content=audit_content,
                auto_process_rules=auto_process_rules,
                audit_data_sources=data_sources_str,
                user_input_key=user_input_key,
                reference_materials=reference_materials_str
            )
            
            state_dict = state.model_dump()
            logger.info(f"调用审核Agent - 字段: {field_name}")
            
            response = await self.client.call_agent(state=state_dict)
            
            # 解析审核结果，取最后一个choice
            if not response.choices:
                return AuditResult(
                    raw_response="",
                    success=False,
                    error="响应中没有choices"
                )
            
            last_choice = response.choices[-1].content
            logger.info(f"审核Agent响应内容: {last_choice}...")
            
            # 解析审核结果
            audit_data = self._parse_audit_result(last_choice)
            
            if audit_data and isinstance(audit_data, dict):
                return AuditResult(
                    audit_fields=audit_data.get("审核字段"),
                    audit_rules=audit_data.get("审核规则"),
                    audit_analysis=audit_data.get("审核分析"),
                    audit_result=audit_data.get("审核结果"),
                    raw_response=last_choice,
                    success=True,
                    error=None
                )
            else:
                error_msg = "无法解析审核结果"
                logger.warning(f"{error_msg}, 原始响应: {last_choice[:200]}...")
                return AuditResult(
                    raw_response=last_choice,
                    success=False,
                    error=error_msg
                )
            
        except Exception as e:
            logger.error(f"字段审核失败: {e}")
            return AuditResult(
                raw_response="",
                success=False,
                error=str(e)
            )
    
    def _format_audit_data_sources(self, data_sources: Union[List[str], Dict[str, Any]]) -> str:
        """
        格式化审核数据源为字符串（只描述数据源，不包含实际内容）
        
        Args:
            data_sources: 审核数据源描述（列表或字典）
            
        Returns:
            str: 格式化的字符串
        """
        try:
            if not data_sources:
                return ""
            
            # 新格式：data_sources 是字符串列表（只描述数据源）
            # 示例：["信息披露申请书 - 标的名称", "房屋所有权证 - 房屋地址", "抵押情况说明 - 全文"]
            if isinstance(data_sources, list):
                return "\n".join(data_sources)
            
            # 兼容旧格式：data_sources 是字典（包含实际数据，已不推荐）
            # {
            #     "材料名称1": {
            #         "字段1": "值1",
            #         "字段2": "值2"
            #     }
            # }
            parts = []
            for material_name, fields in data_sources.items():
                if isinstance(fields, dict):
                    for field_name in fields.keys():
                        parts.append(f"{material_name} - {field_name}")
                elif isinstance(fields, str):
                    # 如果是全文标识
                    parts.append(f"{material_name} - 全文")
            
            return "\n".join(parts)
            
        except Exception as e:
            logger.error(f"格式化审核数据源失败: {e}")
            return ""
    
    def _format_reference_materials(self, materials: Dict[str, Any]) -> str:
        """
        格式化参考材料为字符串
        
        Args:
            materials: 参考材料（结构化对象）
            
        Returns:
            str: 格式化的字符串
        """
        try:
            if not materials:
                return ""
            
            # materials 格式示例：
            # {
            #     "材料名称1 - 字段1": "值1",
            #     "材料名称2 - 字段2": "值2"
            # }
            
            parts = []
            for key, value in materials.items():
                parts.append(f"{key}：{value}")
            
            return "\n".join(parts)
            
        except Exception as e:
            logger.error(f"格式化参考材料失败: {e}")
            return json.dumps(materials, ensure_ascii=False)
    
    def _parse_audit_result(self, content: str) -> Optional[Dict[str, Any]]:
        """
        解析审核结果
        
        从响应内容中提取JSON对象中的审核结果
        
        Args:
            content: 响应内容
            
        Returns:
            Optional[Dict]: 审核数据，解析失败返回None
        """
        try:
            # 尝试匹配 ```json ... ``` 格式
            json_pattern = r'```json\s*(\{.*?\})\s*```'
            match = re.search(json_pattern, content, re.DOTALL)
            
            if match:
                json_str = match.group(1)
                data = json.loads(json_str)
                
                logger.debug(f"解析到的审核JSON数据: {data}")
                return data
            
            # 如果没有匹配到，尝试直接解析整个内容
            try:
                data = json.loads(content)
                logger.debug(f"直接解析审核JSON数据: {data}")
                return data
            except:
                pass
            
            logger.warning(f"无法解析审核结果，原始内容: {content[:200]}...")
            return None
            
        except Exception as e:
            logger.error(f"解析审核结果异常: {e}", exc_info=True)
            return None


class ClassificationAgent:
    """文件分类Agent"""
    
    def __init__(self, client: LLMClient):
        """
        初始化分类Agent
        
        Args:
            client: LLM客户端
        """
        self.client = client
    
    async def classify_file(
        self,
        file_name: str,
        file_content: str,
        types_list: str
    ) -> ClassificationResult:
        """
        分类文件
        
        Args:
            file_name: 文件名称
            file_content: 文件内容（MinerU解析后的Markdown格式）
            types_list: 可选的文件类型列表（格式化的字符串）
            
        Returns:
            ClassificationResult: 分类结果
        """
        try:
            state = ClassificationState(
                file_name=file_name,
                file_content=file_content,
                types_list=types_list
            )
            
            state_dict = state.model_dump()
            logger.info(f"调用分类Agent - 文件: {file_name}")
            
            response = await self.client.call_agent(state=state_dict)
            
            # 解析分类结果，取最后一个choice
            if not response.choices:
                logger.warning(f"分类Agent响应中没有choices")
                return ClassificationResult(
                    raw_response="",
                    success=False,
                    error="响应中没有choices"
                )
            
            last_choice = response.choices[-1].content
            logger.info(f"分类Agent响应内容: {last_choice}")
            
            # 解析分类结果
            classification_data = self._parse_classification_json(last_choice)
            
            if classification_data:
                classified_file_type = classification_data.get('classified_file_type')
                confidence = classification_data.get('confidence')
                reason = classification_data.get('reason')
                
                return ClassificationResult(
                    classified_file_type=classified_file_type,
                    confidence=confidence,
                    reason=reason,
                    raw_response=last_choice,
                    success=classified_file_type is not None,
                    error=None if classified_file_type else "无法解析分类结果"
                )
            else:
                logger.warning(f"无法解析分类结果，完整响应: {last_choice}")
                return ClassificationResult(
                    raw_response=last_choice,
                    success=False,
                    error="无法解析分类结果"
                )
            
        except Exception as e:
            logger.error(f"文件分类失败: {e}", exc_info=True)
            return ClassificationResult(
                raw_response="",
                success=False,
                error=str(e)
            )
    
    def _parse_classification_json(self, content: str) -> Optional[Dict[str, Any]]:
        """
        解析分类结果的 JSON 数据
        
        从响应内容中提取JSON对象
        
        Args:
            content: 响应内容
            
        Returns:
            Optional[Dict]: 包含 classified_file_type, confidence, reason 的字典，解析失败返回None
        """
        try:
            # 尝试匹配 ```json ... ``` 格式
            json_pattern = r'```json\s*(\{.*?\})\s*```'
            match = re.search(json_pattern, content, re.DOTALL)
            
            data = None
            if match:
                json_str = match.group(1)
                data = json.loads(json_str)
            else:
                # 如果没有匹配到，尝试直接解析整个内容
                try:
                    data = json.loads(content)
                except:
                    pass
            
            if data and isinstance(data, dict):
                return {
                    'classified_file_type': data.get('classified_file_type'),
                    'confidence': data.get('confidence'),
                    'reason': data.get('reason')
                }
            
            logger.warning(f"无法解析分类结果，原始内容: {content[:200]}...")
            return None
            
        except Exception as e:
            logger.error(f"解析分类结果异常: {e}")
            return None

